package demo02;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.SocketChannel;
import java.util.Iterator;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

public class TestClient {

    private static final String HOST = "localhost";
    private static final int IO_PORT = 8080;
    private static final int NIO_PORT = 8081;
    private static final String HELLO_MESSAGE = "Hello, Hanserwei!";

    public static void main(String[] args) throws InterruptedException, IOException {
        // 为了方便测试，将客户端数量设置为 10000，如果你需要，可以增加
        int clientCount = 10000;

        // ------------------------- 传统 IO 测试 -------------------------
        System.out.println("开始测试传统 IO...");
        // 传统 IO 客户端使用线程池模拟，每个线程处理一个客户端
        ExecutorService ioExecutorService = Executors.newFixedThreadPool(100);
        long startTimeIO = System.nanoTime();

        for (int i = 0; i < clientCount; i++) {
            ioExecutorService.execute(() -> {
                try (Socket socket = new Socket(HOST, IO_PORT)) {
                    OutputStream out = socket.getOutputStream();
                    InputStream in = socket.getInputStream();
                    out.write(HELLO_MESSAGE.getBytes());
                    byte[] buffer = new byte[1024];
                    in.read(buffer);
                } catch (IOException e) {
                    // 忽略异常，以便快速完成测试
                }
            });
        }
        ioExecutorService.shutdown();
        ioExecutorService.awaitTermination(30, TimeUnit.SECONDS);

        long endTimeIO = System.nanoTime();
        System.out.println("传统 IO 服务器处理 " + clientCount + " 个客户端耗时: " + (endTimeIO - startTimeIO) / 1_000_000 + "ms");
        System.out.println("---------------------------------");


        // ------------------------- NIO 测试 -------------------------
        System.out.println("开始测试 NIO...");
        long startTimeNIO = System.nanoTime();

        // NIO 客户端使用 Selector，在一个线程中管理所有连接
        Selector selector = Selector.open();
        for (int i = 0; i < clientCount; i++) {
            SocketChannel socketChannel = SocketChannel.open();
            socketChannel.configureBlocking(false);
            socketChannel.connect(new InetSocketAddress(HOST, NIO_PORT));
            socketChannel.register(selector, SelectionKey.OP_CONNECT);
        }

        int completedClients = 0;
        ByteBuffer buffer = ByteBuffer.wrap(HELLO_MESSAGE.getBytes());

        while (completedClients < clientCount) {
            selector.select();
            Iterator<SelectionKey> iterator = selector.selectedKeys().iterator();

            while (iterator.hasNext()) {
                SelectionKey key = iterator.next();
                iterator.remove();

                SocketChannel client = (SocketChannel) key.channel();

                if (key.isConnectable()) {
                    if (client.isConnectionPending()) {
                        client.finishConnect();
                    }
                    client.register(selector, SelectionKey.OP_WRITE, buffer.duplicate());
                } else if (key.isWritable()) {
                    ByteBuffer writeBuffer = (ByteBuffer) key.attachment();
                    client.write(writeBuffer);
                    if (!writeBuffer.hasRemaining()) {
                        client.register(selector, SelectionKey.OP_READ);
                    }
                } else if (key.isReadable()) {
                    client.read(ByteBuffer.allocate(1024));
                    client.close();
                    completedClients++;
                }
            }
        }

        long endTimeNIO = System.nanoTime();
        System.out.println("NIO 服务器处理 " + clientCount + " 个客户端耗时: " + (endTimeNIO - startTimeNIO) / 1_000_000 + "ms");
    }
}